# -*- coding: utf-8 -*-
# @Time    : 2019/3/28 23:21
# @Author  : seeledu
# @email   : seeledu@bug.moe
# @File    : BLSTM_bert_att1.py
# @Software: PyCharm
"""
双向LSTM，最后接一个全连接层
每一种参数组合进行 5折交叉验证，使用早期停止法
每一折在验证集上结果最好的模型保存到 result_and_model/intent/版本号/number/best文件夹
"""
from __future__ import print_function
import os
import numpy as np
from keras.callbacks import EarlyStopping, ModelCheckpoint
from keras.layers import Dense, Input, Bidirectional, Embedding, Dropout, LSTM
from keras.layers import TimeDistributed, GlobalAveragePooling1D
from keras.models import Model
from keras.optimizers import Adam, rmsprop
from keras.utils import np_utils
# from gensim.models import KeyedVectors
import keras.backend as K
from keras.preprocessing import sequence
import keras.backend.tensorflow_backend as KTF
import tensorflow as tf

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
# 自适应显存
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config)
KTF.set_session(session)

num_fold = 5  # 5折
sequence_length = 30  # 句子最大长度
embedding_dim = 768  # 词向量长度 AI-lab -200 10G-dict-400 bert 768
epochs = 5000  # 最大迭次数
batch_sizes = [20]  # 一次迭代的例子数
lstm_cells = [50, 100]
dense_cells = [256, 512]  # 隐藏层的单元数
drops = [0.3, 0.5]  # dropout的大小
learning_rate = [1e-3]
num_intent = 24  # Intent类别数量
number = 1000  # 每种参数组合对应一个index
# weight_path = '../embedding/smp/10fold_all_text_without_label_lower_npy/'
data_path = r'../data/npy/'
version = 'BLSTM_BERT'
model_root_path = '../result_and_model/intent/' + version + '/'
if not os.path.exists(model_root_path):
    os.makedirs(model_root_path)
output_file = open(model_root_path + version + ".csv", "w", encoding='utf-8')  # 结果文件
output_file.write("number,fold,lstm_cell,dense_cell,dropout,val_acc,test_acc\n")

print('Loading data......')
test_x = np.load(data_path + 'test_x.npy')
test_x = sequence.pad_sequences(test_x, maxlen=sequence_length)
test_y = np.load(data_path + 'test_intent.npy')
test_y = np_utils.to_categorical(test_y, num_classes=num_intent)

for drop in drops:
    for batch_size in batch_sizes:
        for lstm_cell in lstm_cells:
            for dense_cell in dense_cells:
                # avg_acc = 0
                avg_val_acc = 0  # 5折交叉验证在验证集上的平均准确率
                avg_test_acc = 0  # 5折交叉验证在测试集上的平均准确率
                # mul = []
                for fold in range(num_fold):
                    save_path = model_root_path + str(number) + '/'
                    if not os.path.exists(save_path):
                        os.makedirs(save_path)

                    x_train = np.load(data_path + 'train_x' + str(fold) + '.npy')
                    x_dev = np.load(data_path + 'dev_x' + str(fold) + '.npy')
                    x_train = sequence.pad_sequences(x_train, maxlen=sequence_length)
                    x_dev = sequence.pad_sequences(x_dev, maxlen=sequence_length)

                    y_train = np.load(data_path + 'train_intent_y' + str(fold) + '.npy')
                    y_dev = np.load(data_path + 'dev_intent_y' + str(fold) + '.npy')
                    y_train = np_utils.to_categorical(y_train, num_intent)
                    y_dev = np_utils.to_categorical(y_dev, num_intent)

                    print("Creating Model...")
                    inputs = Input(dtype='float32', shape=(sequence_length, embedding_dim))
                    lstm_1 = Bidirectional(LSTM(lstm_cell))(inputs)

                    dropout = Dropout(drop)(lstm_1)
                    output = Dense(units=num_intent, activation='softmax')(dropout)

                    model = Model(inputs=inputs, outputs=output)

                    # verbose 信息展示模式，0或1
                    # save_best_only 设置为true时，只保存在验证集上性能最好的模型
                    # mode min max auto之一，由检测值monitor而定 min-loss max-val_acc
                    checkpoint = ModelCheckpoint(save_path + str(fold) + 'weights.{epoch:04d}-{val_acc:.4f}.hdf5',
                                                 monitor='val_acc',
                                                 verbose=1,
                                                 save_best_only=True, mode='auto')
                    adam = Adam(lr=learning_rate[0], beta_1=0.9, beta_2=0.999, epsilon=1e-08)

                    model.compile(optimizer=adam, loss='categorical_crossentropy', metrics=['accuracy'])

                    print("Traning Model...")
                    model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs,
                              verbose=1, callbacks=[checkpoint,
                                                    EarlyStopping(
                                                        monitor='val_acc',
                                                        verbose=1,
                                                        patience=20,
                                                        mode='max')],
                              validation_data=(x_dev, y_dev))

                    # 已经训练好的一折的模型按照修改时间排序，最新的就是最好的模型
                    l = os.listdir(path=save_path)
                    # os.path.getmtime返回最近文件修改时间
                    # if not os.path.isdir放置目录下有文件夹
                    l.sort(key=lambda fn: os.path.getmtime(save_path + "/" + fn) if not os.path.isdir(save_path + "/" + fn) else 0)

                    # 在验证集上结果最好的模型 加载并复制到best文件夹
                    model.load_weights(save_path + l[-1])
                    if not os.path.exists(save_path + '/best/'):
                        os.makedirs(save_path + '/best/')
                    model.save(save_path + '/best/' + l[-1])

                    # if len(mul) == 0:
                    #     mul = model.predict(test_x)
                    # else:
                    #     mul += model.predict(test_x)

                    # Returns the loss value & metrics values for the model in test mode.
                    dev_result = model.evaluate(x_dev, y_dev)
                    test_result = model.evaluate(test_x, test_y)
                    avg_val_acc += dev_result[1]
                    avg_test_acc += test_result[1]
                    output_file.write(str(number)+","+str(fold)+","+str(lstm_cell)+","+str(dense_cell)+","
                                      +str(drop)+","+str(dev_result[1])+","+str(test_result[1])+"\n")
                    K.clear_session()
                    del model

                # right = 0
                # for index, tmp in enumerate(mul):
                #     if tmp.argmax() == test_y[index].argmax():
                #         right += 1
                # output_file.write(str(avg_val_acc / 5) + '\t' + str(1.0 * right / mul.shape[0]) + '\n')
                print("Number:", number, "Average Validation accuracy:", str(avg_val_acc / 5), "Average Test accuracy:", str(avg_test_acc / 5))
                # 将缓冲区中的数据立刻写入文件，同时清空缓冲区，不需要是被动的等待输出缓冲区写入
                # output_file.flush()
                # 运行完一种参数组合的5折交叉验证
                number += 1

output_file.close()
